#!/usr/bin/env python
# -*- coding: utf-8 -*-
import mxnet as mx
import logging
import numpy as np
from skimage import io, transform
import cv2
from scipy import misc

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

# Load the pre-trained model
prefix = 'scene_model/scene'
num_round = 20
model = mx.model.FeedForward.load(prefix, num_round, ctx=mx.cpu(), numpy_batch_size=1)
# mean image
mean_rgb = np.array([123.68,116.779,103.939])
mean_rgb = mean_rgb.reshape((3, 1, 1))

# load synset (text label)
synset = [l.strip() for l in open('scene_model/synset.txt').readlines()]
image_path = '6.jpg'

def PreprocessImage(image_path):
    # load image
    rgb_image = io.imread(image_path)
    # resize to 224, 224
    norm_image = transform.resize(rgb_image, (224, 224))
    print type(norm_image)
#    # convert to numpy.ndarray
    sample = np.asarray(norm_image) * 256
    # swap axes to make image from (224, 224, 3) to (3, 224, 224)
    sample = np.swapaxes(sample, 0, 2)
    sample = np.swapaxes(sample, 1, 2)
    # sub mean 
    normed_img = sample - mean_rgb
    normed_img = normed_img.reshape((1, 3, 224, 224))
    return normed_img
    
# Get preprocessed batch (single image batch)
batch = PreprocessImage(image_path)
# Get prediction probability of 1000 classes from model
prob = model.predict(batch)[0]
# Argsort, get prediction index from largest prob to lowest
pred = np.argsort(prob)[::-1]
# Get top1 label
top1 = synset[pred[0]]
print("Top1: ", top1)
# Get top5 label
top5 = [synset[pred[i]] for i in range(5)]
print("Top5: ", top5)
